// Copyright Epic Games, Inc. All Rights Reserved.

#include "up_cmd.h"

#include <zencore/compactbinary.h>
#include <zencore/filesystem.h>
#include <zencore/fmtutils.h>
#include <zencore/logging.h>
#include <zencore/process.h>
#include <zencore/timer.h>
#include <zenutil/zenserverprocess.h>

#include <memory>

namespace zen {

UpCommand::UpCommand()
{
	m_Options.add_option("", "p", "port", "Host port", cxxopts::value(m_Port)->default_value("0"), "<hostport>");
	m_Options.add_option("lifetime",
						 "",
						 "owner-pid",
						 "Specify owning process id",
						 cxxopts::value(m_OwnerPid)->default_value("0"),
						 "<identifier>");
	m_Options.add_option("", "b", "base-dir", "Parent folder of server executable", cxxopts::value(m_ProgramBaseDir), "<directory>");
}

UpCommand::~UpCommand() = default;

int
UpCommand::Run(const ZenCliOptions& GlobalOptions, int argc, char** argv)
{
	ZEN_UNUSED(GlobalOptions);

	if (!ParseOptions(argc, argv))
	{
		return 0;
	}

	{
		ZenServerState State;
		if (State.InitializeReadOnly())
		{
			struct EntryInfo
			{
				uint32_t Pid		   = 0;
				uint16_t DesiredPort   = 0;
				uint16_t EffectivePort = 0;
			};
			std::vector<EntryInfo> RunningEntries;
			State.Snapshot([&RunningEntries, DesiredPort = this->m_Port](const zen::ZenServerState::ZenServerEntry& Entry) {
				if (DesiredPort == 0 || Entry.DesiredListenPort.load() == DesiredPort)
				{
					RunningEntries.push_back(EntryInfo{.Pid			  = Entry.Pid.load(),
													   .DesiredPort	  = Entry.DesiredListenPort.load(),
													   .EffectivePort = Entry.EffectiveListenPort.load()});
				}
			});
			if (RunningEntries.size() > 0)
			{
				ZEN_CONSOLE("Zen server already running with base port {}. First instance at port {}, pid {}",
							RunningEntries[0].DesiredPort,
							RunningEntries[0].EffectivePort,
							RunningEntries[0].Pid);
				return 0;
			}
		}
	}

	if (m_ProgramBaseDir.empty())
	{
		std::filesystem::path ExePath = zen::GetRunningExecutablePath();
		m_ProgramBaseDir			  = ExePath.parent_path();
	}
	ZenServerEnvironment ServerEnvironment;
	ServerEnvironment.Initialize(m_ProgramBaseDir);
	ZenServerInstance Server(ServerEnvironment);
	if (m_OwnerPid != 0)
	{
		Server.SetOwnerPid(m_OwnerPid);
	}
	Server.SpawnServer(m_Port, GlobalOptions.PassthroughCommandLine);

	int Timeout = 10000;

	if (!Server.WaitUntilReady(Timeout))
	{
		if (Server.IsRunning())
		{
			ZEN_ERROR("zen server launch failed (timed out), terminating");
			Server.Terminate();
			return 111;
		}
		int ReturnCode = Server.Shutdown();
		ZEN_CONSOLE("{}", Server.GetLogOutput());
		return ReturnCode;
	}
	else
	{
		ZEN_CONSOLE("zen server up");
	}
	return 0;
}

//////////////////////////////////////////////////////////////////////////

AttachCommand::AttachCommand()
{
	m_Options.add_option("", "p", "port", "Host port", cxxopts::value(m_Port)->default_value("8558"), "<hostport>");
	m_Options.add_option("lifetime", "", "owner-pid", "Specify owning process id", cxxopts::value(m_OwnerPid), "<identifier>");
	m_Options.add_option("", "", "data-dir", "Path to data directory to inspect for running server", cxxopts::value(m_DataDir), "<file>");
}

AttachCommand::~AttachCommand() = default;

int
AttachCommand::Run(const ZenCliOptions& GlobalOptions, int argc, char** argv)
{
	ZEN_UNUSED(GlobalOptions);

	if (!ParseOptions(argc, argv))
	{
		return 0;
	}

	ZenServerState Instance;
	Instance.Initialize();
	Instance.Sweep();
	ZenServerState::ZenServerEntry* Entry = Instance.Lookup(m_Port);

	if (!m_DataDir.empty())
	{
		if (!std::filesystem::is_regular_file(m_DataDir / ".lock"))
		{
			ZEN_CONSOLE("lock file does not exist in directory '{}'", m_DataDir);
			return 1;
		}
		LockFileInfo Info = ReadLockFilePayload(LoadCompactBinaryObject(IoBufferBuilder::MakeFromFile(m_DataDir / ".lock")));
		std::string	 Reason;
		if (!ValidateLockFileInfo(Info, Reason))
		{
			ZEN_CONSOLE("lock file in directory '{}' is not valid. Reason: '{}'", m_DataDir, Reason);
			return 1;
		}
		Entry = Instance.LookupByEffectivePort(Info.EffectiveListenPort);
	}

	if (!Entry)
	{
		ZEN_WARN("no zen server instance to add sponsor process to");
		return 1;
	}

	if (!Entry->AddSponsorProcess(m_OwnerPid, 2000))
	{
		ZEN_WARN("unable to add sponsor process to running zen server instance");
		return 1;
	}

	ZEN_CONSOLE("added sponsor process {} to running instance {} on port {}", m_OwnerPid, Entry->Pid.load(), m_Port);
	return 0;
}

//////////////////////////////////////////////////////////////////////////

DownCommand::DownCommand()
{
	m_Options.add_option("", "p", "port", "Host port", cxxopts::value(m_Port)->default_value("0"), "<hostport>");
	m_Options.add_option("", "f", "force", "Force terminate if graceful shutdown fails", cxxopts::value(m_ForceTerminate), "<force>");
	m_Options.add_option("", "b", "base-dir", "Parent folder of server executable", cxxopts::value(m_ProgramBaseDir), "<directory>");
	m_Options.add_option("", "", "data-dir", "Path to data directory to inspect for running server", cxxopts::value(m_DataDir), "<file>");
}

DownCommand::~DownCommand() = default;

int
DownCommand::Run(const ZenCliOptions& GlobalOptions, int argc, char** argv)
{
	ZEN_UNUSED(GlobalOptions);

	if (!ParseOptions(argc, argv))
	{
		return 0;
	}

	// Discover executing instances
	ZenServerState Instance;
	Instance.Initialize();
	ZenServerState::ZenServerEntry* Entry = Instance.Lookup(m_Port);

	if (m_ProgramBaseDir.empty())
	{
		std::filesystem::path ExePath = zen::GetRunningExecutablePath();
		m_ProgramBaseDir			  = ExePath.parent_path();
	}

	if (!m_DataDir.empty())
	{
		if (!std::filesystem::is_regular_file(m_DataDir / ".lock"))
		{
			ZEN_CONSOLE("lock file does not exist in directory '{}'", m_DataDir);
			return 1;
		}
		LockFileInfo Info = ReadLockFilePayload(LoadCompactBinaryObject(IoBufferBuilder::MakeFromFile(m_DataDir / ".lock")));
		std::string	 Reason;
		if (!ValidateLockFileInfo(Info, Reason))
		{
			ZEN_CONSOLE("lock file in directory '{}' is not valid. Reason: '{}'", m_DataDir, Reason);
			return 1;
		}
		Entry = Instance.LookupByEffectivePort(Info.EffectiveListenPort);
	}

	if (Entry)
	{
		int			   EntryPort		= (int)Entry->DesiredListenPort.load();
		const uint32_t ServerProcessPid = Entry->Pid.load();
		try
		{
			ZenServerEnvironment ServerEnvironment;
			ServerEnvironment.Initialize(m_ProgramBaseDir);
			ZenServerInstance Server(ServerEnvironment);
			Server.AttachToRunningServer(EntryPort);

			ZEN_CONSOLE("attached to server on port {} (pid {}), requesting shutdown", EntryPort, ServerProcessPid);

			Server.Shutdown();

			ZEN_CONSOLE("shutdown complete");

			return 0;
		}
		catch (const std::exception& Ex)
		{
			ZEN_DEBUG("Exception caught when requesting shutdown: {}", Ex.what());
		}

		// Since we cannot obtain a handle to the process we are unable to block on the process
		// handle to determine when the server has shut down. Thus we signal that we would like
		// a shutdown via the shutdown flag and the check if the entry is still running.

		ZEN_CONSOLE("requesting shutdown of server on port {}", EntryPort);
		Entry->SignalShutdownRequest();

		Stopwatch Timer;
		while (Timer.GetElapsedTimeMs() < 5000)
		{
			Instance.Sweep();
			Entry = Instance.Lookup(EntryPort);
			if (Entry == nullptr)
			{
				ZEN_CONSOLE("shutdown complete");
				return 0;
			}
			if (Entry->Pid.load() != ServerProcessPid)
			{
				ZEN_CONSOLE("shutdown complete");
				return 0;
			}
			Sleep(100);
		}
	}
	if (m_ForceTerminate)
	{
		// Try to find the running executable by path name
		std::filesystem::path ServerExePath = m_ProgramBaseDir / "zenserver" ZEN_EXE_SUFFIX_LITERAL;
		ProcessHandle		  RunningProcess;
		if (std::error_code Ec = FindProcess(ServerExePath, RunningProcess); !Ec)
		{
			ZEN_WARN("attempting hard terminate of zen process with pid ({})", RunningProcess.Pid());
			try
			{
				if (RunningProcess.Terminate(0))
				{
					ZEN_CONSOLE("terminate complete");
					return 0;
				}
				ZEN_CONSOLE("failed to terminate server, still running");
				return 1;
			}
			catch (const std::exception& Ex)
			{
				ZEN_CONSOLE("failed to terminate server: '{}'", Ex.what());
				return 1;
			}
		}
		else
		{
			ZEN_CONSOLE("Failed to find process '{}', reason: {}", ServerExePath.string(), Ec.message());
		}
	}
	else if (Entry)
	{
		ZEN_CONSOLE("failed to shutdown of server on port {}, use --force to hard terminate process", Entry->DesiredListenPort.load());
		return 1;
	}

	ZEN_CONSOLE("no zen server to bring down");
	return 0;
}

}  // namespace zen
